import os
from chromadb import HttpClient
from chromadb.utils import embedding_functions
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader, TextLoader, Docx2txtLoader
from typing import List, Dict, Any
import logging

class KnowledgeBase:
    def __init__(self):
        self.chroma_client = None
        self.collection = None
        self.embedding_model = "all-MiniLM-L6-v2"
        self.text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000,
            chunk_overlap=200
        )
        self._init_chroma()

    def _init_chroma(self):
        """Initialize ChromaDB connection"""
        chroma_url = os.getenv("CHROMA_URL", "http://localhost:8000")
        try:
            self.chroma_client = HttpClient(host=chroma_url)
            logging.info(f"成功连接到ChromaDB: {chroma_url}")
        except Exception as e:
            logging.error(f"ChromaDB连接失败: {e}")
            raise
        
        # 创建或获取集合
        embedding_func = embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name=self.embedding_model
        )
        self.collection = self.chroma_client.get_or_create_collection(
            name="financial_docs",
            embedding_function=embedding_func
        )

    def add_document(self, file_path: str):
        """添加文档到知识库"""
        try:
            # 根据文件类型选择加载器
            if file_path.endswith('.pdf'):
                loader = PyPDFLoader(file_path)
            elif file_path.endswith('.docx'):
                loader = Docx2txtLoader(file_path)
            else:
                loader = TextLoader(file_path)
                
            documents = loader.load()
            chunks = self.text_splitter.split_documents(documents)
            
            # 准备文档数据
            ids = [f"doc_{i}" for i in range(len(chunks))]
            texts = [chunk.page_content for chunk in chunks]
            metadatas = [{"source": file_path} for _ in chunks]
            
            # 添加到集合
            self.collection.add(
                ids=ids,
                documents=texts,
                metadatas=metadatas
            )
            logging.info(f"已添加文档: {file_path} (共{len(chunks)}个片段)")
            return True
        except Exception as e:
            logging.error(f"添加文档失败: {file_path}, 错误: {e}")
            return False

    def search(self, query: str, n_results: int = 3) -> List[Dict[str, Any]]:
        """在知识库中搜索相关内容"""
        try:
            results = self.collection.query(
                query_texts=[query],
                n_results=n_results
            )
            return [
                {
                    "document": doc,
                    "metadata": meta,
                    "distance": dist
                }
                for doc, meta, dist in zip(
                    results['documents'][0],
                    results['metadatas'][0],
                    results['distances'][0]
                )
            ]
        except Exception as e:
            logging.error(f"知识库搜索失败: {query}, 错误: {e}")
            return []